# --------------------------------------------------------
# Swin Transformer
# Copyright (c) 2021 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ze Liu
# --------------------------------------------------------'

import os
import yaml
from yacs.config import CfgNode as CN

_C = CN()

# Base config files 基础配置文件的路径
_C.BASE = ['']

# -----------------------------------------------------------------------------
# Data settings  数据配置项
# -----------------------------------------------------------------------------
# 定义一个命名空间（CfgNode）来存储与数据相关的配置
_C.DATA = CN()
# Batch size for a single GPU, could be overwritten by command line argument
# 每个 GPU 上的批量大小,可通过命令行覆盖
_C.DATA.BATCH_SIZE = 128
# Path to dataset, could be overwritten by command line argument
# 数据集路径,可通过命令行覆盖
_C.DATA.DATA_PATH = ''
# Dataset name
# 数据集名称
_C.DATA.DATASET = 'imagenet'
# Input image size
# 输入图像的大小
_C.DATA.IMG_SIZE = 224
# Interpolation to resize image (random, bilinear, bicubic)
# 用于图像缩放的插值方法
_C.DATA.INTERPOLATION = 'bicubic'
# Use zipped dataset instead of folder dataset
# could be overwritten by command line argument
# 是否使用压缩数据集
_C.DATA.ZIP_MODE = False
# Cache Data in Memory, could be overwritten by command line argument
# 数据缓存模式
_C.DATA.CACHE_MODE = 'part'
# Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.
# 是否在 DataLoader 中固定 CPU 内存以加速数据传输到 GPU
_C.DATA.PIN_MEMORY = True
# 数据加载线程的数量
# Number of data loading threads
_C.DATA.NUM_WORKERS = 8

# -----------------------------------------------------------------------------
# Model settings  模型配置项
# -----------------------------------------------------------------------------
# 定义一个命名空间（CfgNode）来存储与模型相关的配置
_C.MODEL = CN()
# Model type  
# 模型类型
_C.MODEL.TYPE = 'swin'
# Model name  
# 模型名称
_C.MODEL.NAME = 'swin_tiny_patch4_window7_224'
# Checkpoint to resume, could be overwritten by command line argument
# 预训练模型的路径
_C.MODEL.PRETRAIN_CKPT = './pretrained_ckpt/swin_tiny_patch4_window7_224.pth'
# 模型训练的恢复路径
_C.MODEL.RESUME = ''
# Number of classes, overwritten in data preparation
# 分类任务的类别数
_C.MODEL.NUM_CLASSES = 1000
# Dropout rate
# Dropout 概率
_C.MODEL.DROP_RATE = 0.0
# Drop path rate
# 随机深度（Stochastic Depth）丢弃的概率
_C.MODEL.DROP_PATH_RATE = 0.1
# Label Smoothing
# 标签平滑的因子
_C.MODEL.LABEL_SMOOTHING = 0.1

# Swin Transformer 的参数设置
# 定义了一个新的配置命名空间，存储与 Swin Transformer 模型相关的参数
_C.MODEL.SWIN = CN()
# Patch 的大小，默认值为 4。图像会被分割成 4x4 的小块，每个块将作为一个输入单元
_C.MODEL.SWIN.PATCH_SIZE = 4
# IN_CHANS：输入图像的通道数，默认值为 3（RGB 图像）
_C.MODEL.SWIN.IN_CHANS = 3
# EMBED_DIM：初始的嵌入维度，默认值为 96
_C.MODEL.SWIN.EMBED_DIM = 96
# DEPTHS：编码器中每个阶段的 Transformer 块数量，默认值 [2, 2, 6, 2]
_C.MODEL.SWIN.DEPTHS = [2, 2, 6, 2]
# DECODER_DEPTHS：解码器中每个阶段的块数量，默认值 [2, 2, 6, 2]
_C.MODEL.SWIN.DECODER_DEPTHS = [2, 2, 6, 2]
# NUM_HEADS：每个阶段的多头注意力头数，默认为 [3, 6, 12, 24]，与嵌入维度成比例
_C.MODEL.SWIN.NUM_HEADS = [3, 6, 12, 24]
# WINDOW_SIZE：Swin Transformer 中窗口的大小，默认值为 7。
_C.MODEL.SWIN.WINDOW_SIZE = 7
# MLP_RATIO：MLP 层的隐藏维度与输入维度的比例，默认值为 4.0。
_C.MODEL.SWIN.MLP_RATIO = 4.
# QKV_BIAS：是否在查询、键和值的线性变换中使用偏置，默认为 True
_C.MODEL.SWIN.QKV_BIAS = True
# QK_SCALE：查询和键的缩放比例，默认为 None。
_C.MODEL.SWIN.QK_SCALE = None
# APE：是否添加绝对位置编码，默认为 False。
_C.MODEL.SWIN.APE = False
# PATCH_NORM：是否对 Patch 进行归一化，默认为 True。
_C.MODEL.SWIN.PATCH_NORM = True
# FINAL_UPSAMPLE：最终的上采样方式，默认为 "expand_first"。
_C.MODEL.SWIN.FINAL_UPSAMPLE= "expand_first"

# -----------------------------------------------------------------------------
# Training settings   训练相关设置
# -----------------------------------------------------------------------------
# 定义了一个新的配置命名空间，存储训练相关的参数。
_C.TRAIN = CN()
# START_EPOCH：训练的起始轮次，默认值为 0。
_C.TRAIN.START_EPOCH = 0
# EPOCHS：训练的总轮次，默认值为 300。
_C.TRAIN.EPOCHS = 300
# WARMUP_EPOCHS：学习率预热的轮次，默认值为 20。
_C.TRAIN.WARMUP_EPOCHS = 20
# WEIGHT_DECAY：权重衰减因子，默认值为 0.05。
_C.TRAIN.WEIGHT_DECAY = 0.05
# BASE_LR：基础学习率，默认值为 5e-4。
_C.TRAIN.BASE_LR = 5e-4
# WARMUP_LR：预热阶段的初始学习率，默认值为 5e-7。
_C.TRAIN.WARMUP_LR = 5e-7
# MIN_LR：最小学习率，默认值为 5e-6。
_C.TRAIN.MIN_LR = 5e-6
# Clip gradient norm
_C.TRAIN.CLIP_GRAD = 5.0
# Auto resume from latest checkpoint
_C.TRAIN.AUTO_RESUME = True
# CLIP_GRAD：梯度裁剪的最大值，用于稳定训练，默认值为 5.0。
_C.TRAIN.ACCUMULATION_STEPS = 0
# Whether to use gradient checkpointing to save memory
# AUTO_RESUME：是否自动从最近的检查点恢复训练，默认值为 True。
_C.TRAIN.USE_CHECKPOINT = False

# LR 学习率调度器
_C.TRAIN.LR_SCHEDULER = CN()
_C.TRAIN.LR_SCHEDULER.NAME = 'cosine'
# Epoch interval to decay LR, used in StepLRScheduler
_C.TRAIN.LR_SCHEDULER.DECAY_EPOCHS = 30
# LR decay rate, used in StepLRScheduler
_C.TRAIN.LR_SCHEDULER.DECAY_RATE = 0.1

# 优化器
_C.TRAIN.OPTIMIZER = CN()
_C.TRAIN.OPTIMIZER.NAME = 'adamw'
# Optimizer Epsilon
_C.TRAIN.OPTIMIZER.EPS = 1e-8
# Optimizer Betas
_C.TRAIN.OPTIMIZER.BETAS = (0.9, 0.999)
# SGD momentum
_C.TRAIN.OPTIMIZER.MOMENTUM = 0.9

# -----------------------------------------------------------------------------
# Augmentation settings  数据增强相关设置
# -----------------------------------------------------------------------------
_C.AUG = CN()
# Color jitter factor 
# 颜色抖动强度，默认值为 0.4
_C.AUG.COLOR_JITTER = 0.4
# Use AutoAugment policy. "v0" or "original"
# 自动增强策略，默认值为 'rand-m9-mstd0.5-inc1'
_C.AUG.AUTO_AUGMENT = 'rand-m9-mstd0.5-inc1'
# Random erase prob
# 随机擦除的概率，默认值为 0.25
_C.AUG.REPROB = 0.25
# Random erase mode
# 随机擦除的模式，默认值为 'pixel'
_C.AUG.REMODE = 'pixel'
# Random erase count
# 随机擦除的次数，默认值为 1
_C.AUG.RECOUNT = 1
# Mixup alpha, mixup enabled if > 0
# Mixup 的 alpha 参数，默认值为 0.8
_C.AUG.MIXUP = 0.8
# Cutmix alpha, cutmix enabled if > 0
# Cutmix 的 alpha 参数，默认值为 1.0。
_C.AUG.CUTMIX = 1.0
# Cutmix min/max ratio, overrides alpha and enables cutmix if set
_C.AUG.CUTMIX_MINMAX = None
# Probability of performing mixup or cutmix when either/both is enabled
# 执行 Mixup 或 Cutmix 的概率，默认值为 1.0。
_C.AUG.MIXUP_PROB = 1.0
# Probability of switching to cutmix when both mixup and cutmix enabled
# 在 Mixup 和 Cutmix 之间切换的概率，默认值为 0.5
_C.AUG.MIXUP_SWITCH_PROB = 0.5
# How to apply mixup/cutmix params. Per "batch", "pair", or "elem"
# Mixup/Cutmix 的应用模式，默认值为 'batch'。
_C.AUG.MIXUP_MODE = 'batch'

# -----------------------------------------------------------------------------
# Testing settings 测试相关设置
# -----------------------------------------------------------------------------
_C.TEST = CN()
# Whether to use center crop when testing  测试时是否使用中心裁剪，默认值为 True。
_C.TEST.CROP = True

# -----------------------------------------------------------------------------
# Misc 杂项 控制混合精度训练的优化等级。
# -----------------------------------------------------------------------------
# Mixed precision opt level, if O0, no amp is used ('O0', 'O1', 'O2')
# 'O0'：不使用混合精度训练。
# 'O1'：使用混合精度训练，并在性能和稳定性之间进行权衡。
# 'O2'：更激进的混合精度优化。
_C.AMP_OPT_LEVEL = ''
# Path to output folder, overwritten by command line argument
# 指定模型训练或评估的输出目录
_C.OUTPUT = ''
# Tag of experiment, overwritten by command line argument
# 实验的标签，用于标识特定的实验配置
_C.TAG = 'default'
# Frequency to save checkpoint
# 保存模型检查点的频率（以 epoch 为单位）
_C.SAVE_FREQ = 1
# Frequency to logging info
# 日志打印的频率（以迭代数为单位）
_C.PRINT_FREQ = 10
# Fixed random seed
# 固定随机种子
_C.SEED = 0
# Perform evaluation only, overwritten by command line argument
# 是否仅执行模型评估
_C.EVAL_MODE = False
# Test throughput only, overwritten by command line argument
# 是否仅测试模型的吞吐量（即每秒处理的样本数量）
_C.THROUGHPUT_MODE = False
# local rank for DistributedDataParallel, given by command line argument
# 分布式训练中本地 GPU 的 rank（ID），默认值0（默认使用第一块 GPU）
_C.LOCAL_RANK = 0

# 递归加载配置文件：从 YAML 文件中加载配置，并与当前配置 config 合并。
# 支持多层嵌套的 BASE 配置：如果 YAML 文件中存在 BASE 键，则会递归加载指定的基础配置文件，并依次合并到主配置中。
def _update_config_from_file(config, cfg_file):
    config.defrost()
    with open(cfg_file, 'r') as f:
        yaml_cfg = yaml.load(f, Loader=yaml.FullLoader)

    for cfg in yaml_cfg.setdefault('BASE', ['']):
        if cfg:
            _update_config_from_file(
                config, os.path.join(os.path.dirname(cfg_file), cfg)
            )
    print('=> merge config from {}'.format(cfg_file))
    config.merge_from_file(cfg_file)
    config.freeze()

# 功能
# 更新配置：将命令行参数 args 中的值合并到 config 中，支持动态覆盖配置。
# 合并优先级：
#   1.从 YAML 文件加载配置（通过 _update_config_from_file）。
#   2.命令行传入的 --opts 参数列表
#   3.命令行直接指定的参数（如 --batch_size、--resume 等）
def update_config(config, args):
    _update_config_from_file(config, args.cfg)

    config.defrost()
    if args.opts:
        config.merge_from_list(args.opts)

    # merge from specific arguments
    if args.batch_size:
        config.DATA.BATCH_SIZE = args.batch_size
    if args.zip:
        config.DATA.ZIP_MODE = True
    if args.cache_mode:
        config.DATA.CACHE_MODE = args.cache_mode
    if args.resume:
        config.MODEL.RESUME = args.resume
    if args.accumulation_steps:
        config.TRAIN.ACCUMULATION_STEPS = args.accumulation_steps
    if args.use_checkpoint:
        config.TRAIN.USE_CHECKPOINT = True
    if args.amp_opt_level:
        config.AMP_OPT_LEVEL = args.amp_opt_level
    if args.tag:
        config.TAG = args.tag
    if args.eval:
        config.EVAL_MODE = True
    if args.throughput:
        config.THROUGHPUT_MODE = True

    config.freeze()

# 功能
#   获取最终的配置对象。
#   默认从全局配置 _C 创建一个副本，并通过 update_config 更新配置。
def get_config(args):
    """Get a yacs CfgNode object with default values."""
    # Return a clone so that the defaults will not be altered
    # This is for the "local variable" use pattern
    config = _C.clone()
    update_config(config, args)

    return config
